#!/usr/bin/env python3
"""
给 Word 里所有内嵌图片统一打水印
用法：
    python word_watermark.py demo.docx "© 2025 公司"
输出文件默认在原文件名后加 _wm.docx
"""
import argparse, io, os, tempfile, shutil, zipfile
from PIL import Image, ImageDraw, ImageFont
from docx import Document
from docx.shared import Cm, Pt

# ---------- 核心：单图加水印 ----------
# opacity:水印透明度
# ---------- 把原 watermark_bytes 整个删掉，贴下面 ----------
import os, platform, io
from PIL import Image, ImageDraw, ImageFont

def get_font_path():
    """返回系统里一定存在的 TTF/OTC 路径，找不到就抛异常"""
    candidates = []
    sys = platform.system()
    if sys == "Windows":
        candidates = [r"C:\Windows\Fonts\msyh.ttc",   # 微软雅黑
                      r"C:\Windows\Fonts\simhei.ttf"] # 黑体
    elif sys == "Darwin":                            # macOS
        candidates = ["/System/Library/Fonts/PingFang.ttc",
                      "/System/Library/Fonts/STHeiti Medium.ttc"]
    else:                                            # Linux
        candidates = ["/usr/share/fonts/opentype/noto/NotoSansCJK-Regular.ttc",
                      "/usr/share/fonts/noto-cjk/NotoSansCJK-Regular.ttc"]
    for p in candidates:
        if os.path.isfile(p):
            return p
    raise FileNotFoundError("找不到中文字体，请手动指定路径")

def watermark_bytes(img_bytes: bytes, text: str, opacity=0.8, angle=30) -> bytes:
    with Image.open(io.BytesIO(img_bytes)) as im:
        im = im.convert("RGBA")
        W, H = im.size

        # 1. 强制 48 px，想再大就改 60/80
        fontsize = 14
        font = ImageFont.truetype(get_font_path(), fontsize)

        # 2. 用 textbbox 算尺寸
        dummy = Image.new("RGBA", (1, 1))
        draw = ImageDraw.Draw(dummy)
        left, top, right, bottom = draw.textbbox((0, 0), text, font=font)
        tw, th = right - left, bottom - top

        # 3. 稀疏平铺：步长放大 4 倍
        txt_layer = Image.new("RGBA", im.size, (255, 255, 255, 0))
        draw = ImageDraw.Draw(txt_layer)
        for y in range(-H, H * 2, int(th * 4)):   # 原来是 2.5，现在 4
            for x in range(-W, W * 2, int(tw * 4)):
                draw.text((x, y), text, font=font,
                          fill=(109, 109, 109, int(255 * opacity)))
        rotated = txt_layer.rotate(angle, expand=0)
        out_im = Image.alpha_composite(im, rotated).convert("RGB")
        buf = io.BytesIO()
        out_im.save(buf, format="PNG", quality=95)
        return buf.getvalue()

# ---------- 主流程 ----------
def process_docx(path: str, text: str, out_path: str = None):
    out_path = out_path or os.path.splitext(path)[0] + "_wm.docx"
    # 用临时目录解压 docx（它本质是 zip）
    with tempfile.TemporaryDirectory() as tmpdir:
        with zipfile.ZipFile(path, 'r') as zf:
            zf.extractall(tmpdir)
        media_dir = os.path.join(tmpdir, "word", "media")
        if not os.path.isdir(media_dir):
            print("未找到内嵌图片，跳过")
            shutil.copy(path, out_path)
            return
        # 遍历所有图片
        changed = False
        for fname in os.listdir(media_dir):
            fpath = os.path.join(media_dir, fname)
            if fname.lower().endswith(('.png', '.jpg', '.jpeg', '.gif', '.bmp')):
                with open(fpath, 'rb') as f:
                    old_bytes = f.read()
                new_bytes = watermark_bytes(old_bytes, text)
                if new_bytes != old_bytes:
                    with open(fpath, 'wb') as f:
                        f.write(new_bytes)
                    changed = True
        if not changed:
            print("图片未变动，直接拷贝")
            shutil.copy(path, out_path)
        else:
            # 重新打包成 docx
            with zipfile.ZipFile(out_path, 'w', zipfile.ZIP_DEFLATED) as zf:
                for folder, _, files in os.walk(tmpdir):
                    for file in files:
                        full = os.path.join(folder, file)
                        arc = os.path.relpath(full, tmpdir)
                        zf.write(full, arc)
        print("已保存：", out_path)

# ---------- 入口 ----------
if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument("docx", help="Word 文件")
    parser.add_argument("text", help="水印文字")
    parser.add_argument("-o", "--output", help="输出文件路径")
    args = parser.parse_args()
    process_docx(args.docx, args.text, args.output)